#!/usr/bin/python3
# -*- coding: utf-8 -*-
"""
@FileName: Tir3_Thomas.py
@Version: 1.0
@Date: 2023/7/20 9:10
@Description: 三对角矩阵，追赶法
"""

import numpy as np
from numpy import ndarray


def thomas(d: ndarray, l: ndarray, u: ndarray, b: ndarray):
    n = len(d)  # 对角线长度
    x = np.zeros(n)
    # 追
    for i in range(1, n):
        l[i - 1] = l[i - 1] / d[i - 1]
        d[i] = d[i] - l[i - 1] * u[i - 1]
        # L.y= b, 回代求解 y
        b[i] = b[i] - b[i - 1] * l[i - 1]
    # 赶; U.x=y 回代求解 x
    x[n - 1] = b[n - 1] / d[n - 1]
    for i in range(n - 2, -1, -1):
        x[i] = (b[i] - u[i] * x[i + 1]) / d[i]
    return x


if __name__ == '__main__':
    d = np.array([13, 14, 15, 16])
    l = np.array([1, 2, 3])
    u = np.array([5, 6, 7])
    b = np.array([2, 4, 6, 7])
    tmp = thomas(d, l, u, b)
    print(tmp)

    mat = np.diag(d, k=0) + np.diag(l, k=-1) + np.diag(u, k=1)
    b2 = mat @ tmp
    print(f'origin b:\n {b}')
    print(f'A.X:\n {b2}')
